"""
用动态规划法解决旅行商问题。
作者：ipid （ https://github.com/ipid ）
"""

from typing import List, Optional, Tuple, cast
from dataclasses import dataclass
from tsp_common_code import read_graph, AdjacencyMatrix


def element_in_set(*, set: int, elem: int) -> bool:
    """
    判断第 element 个元素在不在二进制集合 set 中。
    :param set: 用二进制位所表示的集合
    :param elem: 元素序数
    :return: 若 element 在集合中，即 set 的第 element 个二进制位为 1，则返回 True
    """
    return (1 << elem) & set != 0


def new_set_without_element(*, set: int, elem: int) -> int:
    """
    从 set 中移除掉第 element 个元素，并返回一个新的集合。
    本函数不会改变原集合。

    :param set: 用二进制位所表示的集合
    :param elem: 元素序数
    :return: 新的集合，除了第 element 个元素不在集合里之外，该集合跟 set 一样
    """
    if element_in_set(set=set, elem=elem):
        return set - (1 << elem)
    else:
        return set


@dataclass
class State:
    # 距离。这个值就是动态规划的函数值
    dist: int

    # 该状态对象对应的城市，即 f[i, V] 中的 i
    city_index: int

    # 该状态是从哪个状态转移而来的
    from_state: Optional['State']


class Solution:
    graph: AdjacencyMatrix

    def __init__(self, graph: AdjacencyMatrix):
        self.graph = graph

    def get_initialized_f_array(self) -> List[List[State]]:
        """
        获得动态规划用的状态数组
        :return: 一个二维数组，长度为 [2 ** city_num][city_num]
        """
        city_num = len(self.graph)
        f: List[List[State]] = []
        for i in range(2 ** city_num):
            f.append([cast(State, None)] * city_num)

        for cities_set in range(2 ** city_num):
            for city in range(city_num):
                f[cities_set][city] = State(dist=0, city_index=city, from_state=None)

        for city in range(1, city_num):
            f[0][city].dist = self.graph[city][0]

        return f

    def solve(self) -> Tuple[int, List[int]]:
        """
        使用动态规划法求解问题。
        :return: 一个元组，值为 (最佳距离, 最佳路线)
        """
        city_num = len(self.graph)
        f = self.get_initialized_f_array()

        for cities_set in range(1, 2 ** city_num):
            if element_in_set(set=cities_set, elem=0):
                # 不需要考虑第 0 个城市在点集中的情况
                continue

            for current_city in range(city_num):
                best_result = None
                best_result_state = None

                for next_city in range(city_num):
                    if element_in_set(set=cities_set, elem=next_city):
                        from_state = f[new_set_without_element(set=cities_set, elem=next_city)][next_city]
                        dist = from_state.dist + self.graph[current_city][next_city]

                        if best_result is None or dist < best_result:
                            best_result = dist
                            best_result_state = from_state

                f[cities_set][current_city].dist = best_result
                f[cities_set][current_city].from_state = best_result_state

                city_set_text = ', '.join(reversed(f'{cities_set:b}'.zfill(city_num)))
                print(f'计算结果: f[点集: {city_set_text}][城市: {current_city}] = {best_result}，'
                      f'下一步前往城市 {best_result_state.city_index}')

        result = f[new_set_without_element(set=2 ** city_num - 1, elem=0)][0]
        route: List[int] = []

        state = result
        while state is not None:
            route.append(state.city_index)
            state = state.from_state

        route.append(0)
        return result.dist, route


def main():
    graph = read_graph()
    if len(graph) <= 1:
        result_dist, result_route = 0, []
    else:
        result_dist, result_route = Solution(graph).solve()

    print(f'\n最佳路线：{result_route}\n最佳路线总长度：{result_dist}\n')


if __name__ == '__main__':
    main()
